package com.mlk.core.component.permission;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.CachingExecutor;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.executor.resultset.ResultSetHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.*;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.Objects;
import java.util.Properties;

/**
 * mybatis插件,帮助实现数据权限
 * 该插件会拦截select语句且在数据权限上下文进行过设置的语句.
 */
@Intercepts({@Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})})
@Slf4j
public class DataPermissionPlugin implements Interceptor {


    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        // 获取代理对象之中的所有的参数...
        Object[] args = invocation.getArgs();
        // ------>  index_ms
        int INDEX_MS = 0;
        // 获取的第一个参数就是对应的sql语句对象
        MappedStatement ms = (MappedStatement) args[INDEX_MS];
        // 获取sql的命令类型
        SqlCommandType sqlCommandType = ms.getSqlCommandType();
        if (sqlCommandType != SqlCommandType.SELECT) {
            // 不是select 查询 ,就直接跳过该插件
            return invocation.proceed();
        }

        // 下面判断是否应该查询数据权限
        if (Objects.isNull(DataPermissonContext.get())) {
            // 没有在上下文之中设置值,直接跳过
            DataPermissonContext.remove();
            return invocation.proceed();
        }

        // 下面直接进行处理
        try {
            // 获取了原始的sql对象
            SqlSource sqlSource = ms.getSqlSource();
            // 获取对应的sql语句
            int INDEX_PARAM = 1;
            BoundSql boundSql = sqlSource.getBoundSql(args[INDEX_PARAM]);
            String sql = boundSql.getSql();
            String newSql = DataPermissionSqlFormat.getInstance(DataPermissonContext.get()).formatSql(sql);
            // 重新组织sql语句
            BoundSql newBoundSql = new BoundSql(ms.getConfiguration(), newSql, boundSql.getParameterMappings(), boundSql.getParameterObject());

            // 把新的查询放到statement里
            MappedStatement newMs = copyFromMappedStatement(ms, new BoundSqlSqlSource(newBoundSql));

            // 获取对应的参数映射
            for (ParameterMapping mapping : boundSql.getParameterMappings()) {
                String prop = mapping.getProperty();
                if (boundSql.hasAdditionalParameter(prop)) {
                    newBoundSql.setAdditionalParameter(prop, boundSql.getAdditionalParameter(prop));
                }
            }
            args[INDEX_MS] = newMs;
            // 重新执行语句
            return invocation.proceed();
        } finally {
            // 移除数据,放置内存溢出
            DataPermissonContext.remove();
        }
    }

    /**
     * 复制原始的sql绑定器的附属信息
     * @param ms
     * @param newSqlSource
     * @return
     */
    private MappedStatement copyFromMappedStatement(MappedStatement ms, SqlSource newSqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), newSqlSource, ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        if (ms.getKeyProperties() != null && ms.getKeyProperties().length > 0) {
            builder.keyProperty(ms.getKeyProperties()[0]);
        }
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.resultSetType(ms.getResultSetType());
        builder.cache(ms.getCache());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }

    public static class BoundSqlSqlSource implements SqlSource {
        private BoundSql boundSql;

        public BoundSqlSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
    }
}